from global_var import *
from dataset import *
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.nn.functional as F

if __name__ == "__main__":
    preds = []
    result = []
    for fold in range(k_folds):
        _exp_name = f"model_fold_{fold}"
        model_best = ResidualNet().to(device)
        state_dict = torch.load(f"{_exp_name}_best.ckpt", map_location=device)
        new_state_dict = {
            k.replace("module.", ""): v for k, v in state_dict.items()
        }
        model_best.load_state_dict(new_state_dict)
        model_best.eval()
        test_preds = []
        test_tfm_preds = []
        fold_preds = []
        with torch.no_grad():
            for data,_ in tqdm(test_loader):
                test_pred = model_best(data.to(device))
                test_preds.append(test_pred)
            for data,_ in tqdm(test_tfm_loader):
                test_pred = model_best(data.to(device))
                test_tfm_preds.append(test_pred)
            for idx in range(len(test_preds)):
                test_pred = 0.8 * test_preds[idx] + 0.2 * test_tfm_preds[idx]
                fold_preds.append(test_pred)
        
        preds.append(fold_preds)
    for i in range(len(preds[0])):
        test_pred = sum([preds[fold][i] for fold in range(k_folds)]) 
        test_label = np.argmax(test_pred.cpu().data.numpy(), axis=1)
        result += test_label.tolist()
    print(len(result))
        
    # create test csv
    def pad4(i):
        return "0"*(4-len(str(i)))+str(i)
    df = pd.DataFrame()
    df["Id"] = [pad4(i + 1) for i in range(len(test_set))]
    df["Category"] = result
    df.to_csv("submission.csv",index = False)

